package cc.shacocloud.luanniao.core.service.text.splitter

import cc.shacocloud.luanniao.core.model.constant.DocumentLanguageEnum
import cc.shacocloud.luanniao.core.service.ai.doubao.DoubaoAiModel
import cc.shacocloud.luanniao.core.service.ai.doubao.DoubaoChatOptions
import cc.shacocloud.luanniao.core.service.ai.qaDocuments
import cc.shacocloud.luanniao.core.service.text.extract.pdf.PdfFileContentExtractServiceImpl
import org.junit.jupiter.api.Test
import org.springframework.util.ResourceUtils

/**
 *
 * @author 思追(shaco)
 */
class RecursiveCharacterTextSplitterServiceImplTest {

    @Test
    fun splitText() {

        val aiModel = DoubaoAiModel(
            "4b771ac9-47c6-4383-9c2a-4d105a54db6c",
            "https://ark.cn-beijing.volces.com/api/v3/",
            "ep-20240628082932-bjv8f",
            DoubaoChatOptions(
                model = "doubao-pro-128k"
            )
        )

        val file = ResourceUtils.getFile("classpath:毛选.pdf")
        val extractService = PdfFileContentExtractServiceImpl()
        val content = extractService.extract(file)

        val textSplitterService = RecursiveCharacterTextSplitterServiceImpl(
            maxTokens = 500,
            chunkOverlap = 50,
            aiModel = aiModel
        )

        val texts = textSplitterService.cleanAndSplitText(content)
        for (text in texts.subList(0, 5)) {
            println(text)
            println()
        }

        println("===========================")

        val qaDocuments = aiModel.qaDocuments(texts[1], DocumentLanguageEnum.CHINESE)
        for (qaDocument in qaDocuments) {
            println(qaDocument)
        }

    }
}